import pandas as pd
import numpy as np
import  matplotlib.pyplot as plt
from statsmodels.graphics.tsaplots import plot_acf
from statsmodels.graphics.tsaplots import plot_pacf
from statsmodels.stats.diagnostic import acorr_ljungbox
from statsmodels.tsa.stattools import adfuller as ADF
from statsmodels.tsa.arima.model import ARIMA as arima



class ARIMAModel:
  def __init__(self, data):
    self.data = data
    self.D_data = None
    self.model = None
    pd.set_option('display.float_format', '{:.25f}'.format) #改变显示精度
    pd.set_option('display.precision', 25)

  @staticmethod
  def tagADF(data):
    '''
    功能：对传入的时间序列进行ADF平稳性检验，输出统计量和置信区间，返回是否平稳（p<0.05）。
    参数：data为一维序列（Series或DataFrame的一列）
    '''
    t = ADF(data, autolag='AIC')
    result = pd.DataFrame(index=[
      "Test Statistic Value", "p-value", "Lags Used", 
      "Number of Observations Used", 
      "Critical Value(1%)", "Critical Value(5%)", "Critical Value(10%)"
    ], columns=['value'])
    result.loc['Test Statistic Value', 'value'] = t[0]
    result.loc['p-value', 'value'] = t[1]
    result.loc['Lags Used', 'value'] = t[2]
    result.loc['Number of Observations Used', 'value'] = t[3]
    for key, value in t[4].items():
      result.loc['Critical Value(%s)' % key, 'value'] = value
    print(result['value'])
    return

  def plot_all(self):
    '''
    功能：绘制原始序列及一阶差分的趋势图、ACF图和PACF图，辅助判断平稳性和定阶。
    '''
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    fig, axs = plt.subplots(2, 3, figsize=(12, 8))
    self.data.plot(ax=axs[0, 0], title='原始序列')
    plot_acf(self.data, ax=axs[0, 1])
    plot_pacf(self.data, ax=axs[0, 2])
    if self.D_data is not None:
      self.D_data.plot(ax=axs[1, 0], title='一阶差分')
      plot_acf(self.D_data, ax=axs[1, 1])
      plot_pacf(self.D_data, ax=axs[1, 2])
    plt.tight_layout()
    plt.show()

  def difference(self):
    '''
    用法：arima_model.difference()
    功能：对原始序列进行一阶差分，消除趋势，使序列趋于平稳。
    返回：一阶差分后的序列（Series）
    '''
    self.D_data = self.data.diff(1).dropna()
    return self.D_data

  def white_noise_test(self):
    '''
    用法：arima_model.white_noise_test()
    功能：对白噪声进行检验，判断差分序列是否为白噪声（残差无自相关）。p值不为0说明不是白噪声
    '''
    print('白噪声检验:', acorr_ljungbox(self.D_data, lags=1))

  def select_order(self):
    '''
    功能：利用BIC准则自动选择ARIMA模型的p、q阶数。
    返回：最优的p、q值
    '''
    if not isinstance(self.D_data.index, pd.DatetimeIndex):
      start_date = pd.Timestamp('2023-01-01')
      dates = pd.date_range(start=start_date, periods=len(self.D_data), freq='D')
      self.D_data.index = dates
    pmax = int(len(self.data) / 10)
    qmax = int(len(self.data) / 10)
    bic_matrix = []
    for p in range(pmax + 1):
      tmp = []
      for q in range(qmax + 1):
        try:
          model = arima(self.data, order=(p, 1, q)).fit()
          tmp.append(model.bic)
        except:
          tmp.append(None)
      bic_matrix.append(tmp)
    bic_matrix = pd.DataFrame(bic_matrix)
    p, q = bic_matrix.stack().idxmin()
    print(f'BIC最小的p值和q值为：{p}、{q}')
    return p, q

  def fit_predict(self, p, q, forecast_steps=5):
    '''
    用法：arima_model.fit_predict(p, q, forecast_steps=5)
    功能：拟合ARIMA模型并进行未来数据预测，输出模型报告和预测结果。
    参数：p、q为模型阶数，forecast_steps为预测步数
    '''
    self.model = arima(self.data, order=(p, 1, q)).fit()
    print(self.model.summary())
    print('未来预测:', self.model.forecast(forecast_steps))


# 示例流程
if __name__ == '__main__':
  np.random.seed(0)
  n = 200
  x = np.cumsum(np.random.randn(n))
  data = pd.DataFrame(x, columns=['value'])
  arima_model = ARIMAModel(data)
  arima_model.plot_all()
  print('原始序列的ADF检验结果:')
  arima_model.tagADF(data)
  D_data = arima_model.difference()
  arima_model.plot_all()
  print('差分序列的ADF检验结果:')
  arima_model.tagADF(D_data)
  arima_model.white_noise_test()
  p, q = arima_model.select_order()
  arima_model.fit_predict(p, q)